import os
from openpyxl import load_workbook
from openpyxl import Workbook
from argparse import ArgumentParser
import lief
from telnet_openocd import telnet_openocd
from test_case import test_case
import datetime

def main():

    parser = ArgumentParser()
    parser.add_argument('-i', '--input', type=str, required=True, help="input test case xlsx")
    parser.add_argument('-s', '--symbol', type=str, required=True, help="input symbol file")
    parser.add_argument('-o', '--output', type=str, required=False, help="output test case report xlsx")
    args = parser.parse_args()
    
    if os.path.exists(args.symbol):
        binary: lief.ELF.Binary = lief.ELF.parse(args.symbol)
        # for s in binary.symbols:
        #     print(s.name)
        #     print(hex(s.value))
        # for f in binary.functions:
        #     print(f.name)
        #     print(hex(f.address))
    else:
        print("input symbol file path error!")
        exit()

    if os.path.exists(args.input):
        input_workbook = load_workbook(args.input)
        output_workbook = Workbook()
    else:
        print("input file path error!")
        exit()

    try:
        tn_session = telnet_openocd()
        if tn_session == None:
            print("openocd server is not set up\n")
            exit()
        tn_session.reset_init()
        tn_session.reset_halt()

        main_address = test_case.find_function_address(binary, 'main')
        if main_address != 0xffffffff:
            tn_session.set_breakpoint(main_address)
            tn_session.resume()

            for sheet_index, worksheet in enumerate(input_workbook.worksheets):
                if sheet_index != 0:
                    output_workbook.create_sheet(worksheet.title)
                else:
                    output_workbook.worksheets[sheet_index].title = worksheet.title
                for row in worksheet.iter_rows(values_only=True):
                    output_workbook.worksheets[sheet_index].append(row)

                testcase_list = [cell.value for cell in worksheet['A']]
                description_list = [cell.value for cell in worksheet['B']]
                initial_condition_list = [cell.value for cell in worksheet['C']]
                check_point_list = [cell.value for cell in worksheet['D']]
                test_result_symbol_list = [cell.value for cell in worksheet['E']]
                test_log_symbol_list = [cell.value for cell in worksheet['F']]
                for index, fut_name in enumerate(testcase_list, start = 0):
                    if index == 0:
                        continue
                    if fut_name != None:
                        tc = test_case(fut_name, 
                                    test_case.find_function_address(binary, fut_name), 
                                    description_list[index], 
                                    test_case.find_function_address(binary, initial_condition_list[index]), 
                                    int(check_point_list[index]),
                                    test_case.find_symbol_address(binary, test_result_symbol_list[index]),
                                    test_case.find_symbol_address(binary, test_log_symbol_list[index]))
                    if tc.fut_address != 0xffffffff:
                        print(f"{fut_name} test begin\n")
                        tc.execute_test_case(tn_session)
                    else:
                        print(f"tese case entry is not found : {fut_name}\n")
                    test_result_string_data = ''.join(filter(str.isprintable, tc.test_result))
                    output_workbook.worksheets[sheet_index].cell(row=index + 1, column=5, value=test_result_string_data)
                    test_log_string_data = ''.join(filter(str.isprintable, tc.test_log))
                    output_workbook.worksheets[sheet_index].cell(row=index + 1, column=6, value=test_log_string_data)
        else:
            print("main address is invalid\n")
            exit()

    except Exception as e:
        print(f"Failed to connect to target: {e}")
        exit()

    if args.output == None:
        now = datetime.datetime.now()
        output_workbook_filename = f"testcase_report_{now.strftime('%Y-%m-%d-%H_%M_%S')}.xlsx"
    else:
        output_workbook_filename = args.output
    print(output_workbook_filename)
    output_workbook.save(filename = output_workbook_filename)
main()



